package com.example.dobs.demo.flink.io.tool;

import com.google.protobuf.Message;
import org.apache.flink.api.common.accumulators.LongCounter;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.api.common.serialization.DeserializationSchema;
import org.apache.flink.api.common.serialization.SerializationSchema;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.typeutils.TypeExtractor;

import java.io.IOException;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;

public class ProtoMessageSerializer implements DeserializationSchema<Message>, SerializationSchema<Message> {

    final private Class<?> pbClass;
    private LongCounter sourceErrCount = new LongCounter();
    private LongCounter sourceInCount = new LongCounter();

    public ProtoMessageSerializer(Class<?> pbClass) throws NoSuchMethodException {
        this.pbClass = pbClass;
    }

    @Override
    public Message deserialize(byte[] bytes) throws IOException {
        sourceInCount.add(1);
        try {
            Message message = (Message) this.pbClass.getMethod("parseFrom", byte[].class).invoke(null, bytes);
//            invokeSuccessCount.inc();
            return message;
        } catch (IllegalAccessException e) {
            e.printStackTrace();
        } catch (InvocationTargetException | NoSuchMethodException e) {
            e.printStackTrace();
        }
//        invokeFailCount.inc();
        sourceErrCount.add(1);
        return null;
    }

    @Override
    public boolean isEndOfStream(Message msg) {
        return false;
    }

    @Override
    public TypeInformation<Message> getProducedType() {
        return TypeExtractor.getForClass(Message.class);
    }

    @Override
    public byte[] serialize(Message msg) {
        return msg.toByteArray();
    }
}
